package com.authsphere.security.token.redis;

import com.authsphere.security.token.common.AccessToken;
import com.authsphere.security.token.common.AuthenticationKeyGenerator;
import com.authsphere.security.token.common.TokenStore;
import com.authsphere.security.token.redis.strategy.JdkSerializationStrategy;
import com.authsphere.security.token.redis.strategy.RedisTokenStoreSerializationStrategy;
import org.springframework.data.redis.connection.RedisConnection;
import org.springframework.data.redis.connection.RedisConnectionFactory;
import org.springframework.data.redis.connection.RedisKeyCommands;
import org.springframework.data.redis.connection.RedisStringCommands;
import org.springframework.security.core.Authentication;

import java.util.List;

/**
 * @program: AuthSphere
 * @description:
 * @author: YuKai Fan
 * @create: 2025/3/12 09:32
 **/
public class RedisTokenStore implements TokenStore {
    private static final String ACCESS = "access:";
    private static final String AUTH_TO_ACCESS = "auth_to_access:";
    private static final String AUTH = "auth:";

    private final RedisConnectionFactory connectionFactory;
    private AuthenticationKeyGenerator authenticationKeyGenerator = new RedisAuthenticationKeyGenerator();
    private RedisTokenStoreSerializationStrategy serializationStrategy = new JdkSerializationStrategy();

    private String prefix = "";
    public void setPrefix(String prefix) {
        this.prefix = prefix;
    }

    public RedisTokenStore(RedisConnectionFactory connectionFactory) {
        this.connectionFactory = connectionFactory;
    }

    public void setAuthenticationKeyGenerator(AuthenticationKeyGenerator authenticationKeyGenerator) {
        this.authenticationKeyGenerator = authenticationKeyGenerator;
    }

    public void setSerializationStrategy(RedisTokenStoreSerializationStrategy serializationStrategy) {
        this.serializationStrategy = serializationStrategy;
    }

    private RedisConnection getConnection() {
        return connectionFactory.getConnection();
    }

    private byte[] serialize(Object object) {
        return serializationStrategy.serialize(object);
    }

    private byte[] serializeKey(String object) {
        return serialize(prefix + object);
    }

    private AccessToken deserializeAccessToken(byte[] bytes) {
        return serializationStrategy.deserialize(bytes, AccessToken.class);
    }

    private Authentication deserializeAuthentication(byte[] bytes) {
        return serializationStrategy.deserialize(bytes, Authentication.class);
    }


    @Override
    public AccessToken getAccessToken(Authentication authentication) {
        String key = authenticationKeyGenerator.extractKey(authentication);
        byte[] serializeKey = serializeKey(AUTH_TO_ACCESS + key);
        RedisConnection conn = getConnection();
        byte[] bytes = null;
        try {
            bytes = conn.stringCommands().get(serializeKey);
        } finally {
            conn.close();
        }
        AccessToken accessToken = deserializeAccessToken(bytes);
        if (accessToken != null) {
            Authentication storedAuthentication = readAuthentication(accessToken);
            if (storedAuthentication == null || !key.equals(authenticationKeyGenerator.extractKey(storedAuthentication))) {
                // Keep the stores consistent (maybe the same user is
                // represented by this authentication but the details have
                // changed)
                storeAccessToken(accessToken, authentication);
            }
        }
        return accessToken;
    }

    @Override
    public Authentication readAuthentication(String token) {
        byte[] bytes = null;
        try (RedisConnection conn = getConnection()) {
            RedisStringCommands redisStringCommands = conn.stringCommands();
            bytes = redisStringCommands.get(serializeKey(AUTH + token));
        }
        return deserializeAuthentication(bytes);
    }

    @Override
    public void storeAccessToken(AccessToken token, Authentication authentication) {
        byte[] serializedAccessToken = serialize(token);
        byte[] serializedAuth = serialize(authentication);
        byte[] accessKey = serializeKey(ACCESS + token.getSeriesId());
        byte[] authKey = serializeKey(AUTH + token.getValue());
        byte[] authToAccessKey = serializeKey(AUTH_TO_ACCESS + authenticationKeyGenerator.extractKey(authentication));
        try (RedisConnection conn = getConnection()) {
            conn.openPipeline();
            RedisStringCommands redisStringCommands = conn.stringCommands();
            redisStringCommands.set(accessKey, serializedAccessToken);
            redisStringCommands.set(authKey, serializedAuth);
            redisStringCommands.set(authToAccessKey, serializedAccessToken);

            if (token instanceof ExpirationAccessToken accessToken) {
                int seconds = accessToken.getExpiresIn();
                if (seconds > 0) {
                    RedisKeyCommands keyCommands = conn.keyCommands();
                    keyCommands.expire(accessKey, seconds);
                    keyCommands.expire(authKey, seconds);
                    keyCommands.expire(authToAccessKey, seconds);
                }
            }
            conn.closePipeline();
        }
    }

    @Override
    public Authentication readAuthentication(AccessToken token) {
        return readAuthentication(token.getValue());
    }

    @Override
    public AccessToken readAccessToken(String tokenValue) {
        byte[] key = serializeKey(ACCESS + tokenValue);
        byte[] bytes = null;
        try (RedisConnection conn = getConnection()) {
            RedisStringCommands redisStringCommands = conn.stringCommands();
            bytes = redisStringCommands.get(key);
        }
        return deserializeAccessToken(bytes);
    }

    @Override
    public void removeAccessToken(AccessToken token) {
        removeAccessToken(token.getValue());
    }

    @Override
    public void removeAccessToken(String tokenValue) {
        byte[] accessKey = serializeKey(ACCESS + tokenValue);
        byte[] authKey = serializeKey(AUTH + tokenValue);
        try (RedisConnection conn = getConnection()) {
            conn.openPipeline();
            RedisStringCommands redisStringCommands = conn.stringCommands();
            RedisKeyCommands redisKeyCommands = conn.keyCommands();
            redisStringCommands.get(accessKey);
            redisStringCommands.get(authKey);
            redisKeyCommands.del(accessKey);
            // Don't remove the refresh token - it's up to the caller to do that
            redisKeyCommands.del(authKey);
            List<Object> results = conn.closePipeline();
            byte[] auth = (byte[]) results.get(1);

            Authentication authentication = deserializeAuthentication(auth);
            if (authentication != null) {
                String key = authenticationKeyGenerator.extractKey(authentication);
                byte[] authToAccessKey = serializeKey(AUTH_TO_ACCESS + key);
                conn.openPipeline();
                redisKeyCommands.del(authToAccessKey);
                redisKeyCommands.del(serialize(ACCESS + key));
                conn.closePipeline();
            }
        }
    }
}
